#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
批量随机替换图片背景脚本
功能：从指定背景图目录随机选择背景，批量替换 ready 目录下的图片
"""

# ⚠️ 重要：必须在所有 import 之前设置环境变量，强制使用 CPU 模式
import os
os.environ['CUDA_VISIBLE_DEVICES'] = '-1'              # 禁用 CUDA
os.environ['ORT_DISABLE_PROVIDER'] = 'DML'             # 禁用 DirectML
os.environ['ONNXRUNTIME_EXECUTION_PROVIDERS'] = 'CPUExecutionProvider'  # 强制使用 CPU

import random
from pathlib import Path
from PIL import Image
from typing import List
import logging
import sys

# 🔍 检查 NumPy 版本兼容性（必须在导入 rembg 之前）
try:
    import numpy as np
    numpy_version = tuple(map(int, np.__version__.split('.')[:2]))
    if numpy_version[0] >= 2:
        print("=" * 70)
        print("❌ 错误：检测到不兼容的 NumPy 版本")
        print(f"   当前版本: {np.__version__}")
        print(f"   需要版本: < 2.0.0")
        print()
        print("📋 解决方案：")
        print("   pip install \"numpy<2.0\"")
        print()
        print("💡 或使用完整的依赖安装：")
        print("   pip install -r requirements_bg_replace.txt")
        print("=" * 70)
        sys.exit(1)
except ImportError:
    print("❌ 错误：未安装 NumPy，请先安装依赖")
    print("   pip install numpy pillow rembg")
    sys.exit(1)

# 导入 rembg 并显式创建 CPU session
try:
    from rembg import remove, new_session
except Exception as e:
    print("=" * 70)
    print("❌ 错误：无法导入 rembg 库")
    print(f"   错误详情: {e}")
    print()
    print("📋 建议的解决方案：")
    print("   1. 检查依赖版本:")
    print("      pip list | findstr \"numpy onnxruntime rembg\"")
    print()
    print("   2. 重新安装兼容版本:")
    print("      pip uninstall numpy onnxruntime rembg -y")
    print("      pip install -r requirements_bg_replace.txt")
    print("=" * 70)
    sys.exit(1)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


class BackgroundReplacer:
    """背景替换处理类 - 支持随机背景替换"""
    
    def __init__(self, background_dir_name: str, source_dir: str, picture_base_dir: str = "/Users/wuyunfan/test/picture"):
        """
        初始化背景替换器
        
        Args:
            background_dir_name: picture 目录下的子目录名称（背景图所在目录）
            source_dir: 源图片目录 (ready目录)
            picture_base_dir: picture 根目录路径
        """
        self.source_dir = Path(source_dir)
        self.background_dir = Path(picture_base_dir) / background_dir_name
        
        # 检查背景图目录是否存在
        if not self.background_dir.exists():
            raise ValueError(f"背景图目录不存在: {self.background_dir}")
        
        # 加载所有背景图片
        self.background_images = self._load_background_images()
        
        if not self.background_images:
            raise ValueError(f"背景图目录下没有找到图片: {self.background_dir}")
        
        logger.info(f"✅ 从 {self.background_dir} 加载了 {len(self.background_images)} 张背景图")
        
        # 🔧 显式创建使用 CPU 的 rembg session（解决 Windows DirectML 错误）
        try:
            logger.info("🔧 正在初始化 CPU 模式的背景去除引擎...")
            self.rembg_session = new_session("u2net", providers=["CPUExecutionProvider"])
            logger.info("✅ CPU 模式初始化成功")
        except Exception as e:
            logger.warning(f"⚠️ 无法创建指定 session，将使用默认配置: {e}")
            self.rembg_session = None
    
    def _load_background_images(self) -> List[Path]:
        """
        加载背景图目录下的所有图片文件
        
        Returns:
            背景图片路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'}
        background_files = []
        
        for file in self.background_dir.iterdir():
            if file.is_file() and file.suffix in image_extensions:
                background_files.append(file)
        
        return background_files
    
    def get_image_files(self) -> List[Path]:
        """
        获取目录下所有图片文件
        
        Returns:
            图片文件路径列表
        """
        image_extensions = {'.jpg', '.jpeg', '.png', '.JPG', '.JPEG', '.PNG'}
        image_files = []
        
        for root, dirs, files in os.walk(self.source_dir):
            for file in files:
                if Path(file).suffix in image_extensions:
                    image_files.append(Path(root) / file)
        
        logger.info(f"找到 {len(image_files)} 个图片文件")
        return image_files
    
    def replace_background(self, input_path: Path) -> bool:
        """
        随机替换单张图片的背景并覆盖原文件
        
        Args:
            input_path: 输入图片路径（将被直接修改）
            
        Returns:
            是否成功
        """
        try:
            logger.info(f"📷 处理图片: {input_path.relative_to(self.source_dir)}")
            
            # 随机选择一张背景图
            random_bg_path = random.choice(self.background_images)
            logger.info(f"   🎲 随机背景: {random_bg_path.name}")
            
            # 读取原图
            input_image = Image.open(input_path)
            
            # 使用 rembg 去除背景，显式指定使用 CPU session
            if self.rembg_session:
                foreground = remove(input_image, session=self.rembg_session)
            else:
                foreground = remove(input_image)
            
            # 读取并调整背景图片尺寸以匹配前景
            background = Image.open(random_bg_path).convert("RGB")
            background_resized = background.resize(foreground.size, Image.Resampling.LANCZOS)
            
            # 创建最终图片（RGB模式）
            final_image = Image.new("RGB", foreground.size, (255, 255, 255))
            final_image.paste(background_resized, (0, 0))
            
            # 将前景合成到背景上（使用 alpha 通道进行混合）
            final_image.paste(foreground, (0, 0), foreground)
            
            # 直接覆盖原文件
            final_image.save(input_path, quality=95)
            logger.info(f"   ✅ 保存成功: {input_path.name}")
            
            return True
            
        except Exception as e:
            logger.error(f"❌ 处理图片 {input_path.name} 失败: {str(e)}")
            return False
    
    def batch_replace(self):
        """批量随机替换背景"""
        image_files = self.get_image_files()
        
        if not image_files:
            logger.warning("⚠️  未找到任何图片文件")
            return
        
        success_count = 0
        fail_count = 0
        
        logger.info("=" * 60)
        logger.info(f"🚀 开始批量处理，共 {len(image_files)} 张图片")
        logger.info("=" * 60)
        
        for idx, image_path in enumerate(image_files, 1):
            logger.info(f"\n📊 进度: {idx}/{len(image_files)}")
            
            if self.replace_background(image_path):
                success_count += 1
            else:
                fail_count += 1
        
        logger.info("\n" + "=" * 60)
        logger.info(f"🎉 处理完成！")
        logger.info(f"✅ 成功: {success_count} 张")
        logger.info(f"❌ 失败: {fail_count} 张")
        logger.info(f"📁 处理目录: {self.source_dir}")
        logger.info("=" * 60)


def main():
    """
    主函数
    
    使用方法：
        python batch_replace_background.py
    
    配置说明：
        - BACKGROUND_DIR_NAME: picture 目录下的子目录名称
        - READY_DIR: ready 目录路径（默认 /Users/wuyunfan/test/ready）
        - PICTURE_BASE_DIR: picture 根目录路径（默认 /Users/wuyunfan/test/picture）
    """
    import sys
    
    # 配置参数
    READY_DIR = "/Users/wuyunfan/test/ready"  # ready 目录
    PICTURE_BASE_DIR = "/Users/wuyunfan/test/picture"  # picture 根目录
    
    # 从命令行参数获取背景图目录名称，或使用默认值
    if len(sys.argv) > 1:
        BACKGROUND_DIR_NAME = sys.argv[1]
    else:
        # 默认背景图目录名称（请修改为实际目录名）
        BACKGROUND_DIR_NAME = "default_background"
        print(f"⚠️  未指定背景图目录，使用默认值: {BACKGROUND_DIR_NAME}")
        print(f"💡 使用方法: python batch_replace_background.py <背景图目录名>")
        print()
    
    print("=" * 60)
    print(f"🎨 批量随机背景替换工具")
    print("=" * 60)
    print(f"📂 背景图目录: {PICTURE_BASE_DIR}/{BACKGROUND_DIR_NAME}")
    print(f"📂 处理目录: {READY_DIR}")
    print(f"⚠️  警告: 将直接修改原文件，请确保已备份！")
    print("=" * 60)
    
    # 确认是否继续
    response = input("\n是否继续？(yes/y 继续，其他退出): ").strip().lower()
    if response not in ['yes', 'y']:
        print("❌ 已取消操作")
        return
    
    try:
        # 创建替换器并执行批量替换
        replacer = BackgroundReplacer(
            background_dir_name=BACKGROUND_DIR_NAME,
            source_dir=READY_DIR,
            picture_base_dir=PICTURE_BASE_DIR
        )
        
        replacer.batch_replace()
        
    except ValueError as e:
        logger.error(f"❌ 配置错误: {e}")
        sys.exit(1)
    except Exception as e:
        logger.error(f"❌ 发生错误: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()
